package com.junjie.checkin.websocket.interceptor;

import com.junjie.checkin.security.service.AuthDetails;
import com.junjie.checkin.security.service.impl.JwtServiceImpl;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.security.Principal;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @author 洪浚杰
 * @description WebSocket握手拦截器，用于处理认证
 * @date 2025/6/23 10:01
 */
@Component
public class AuthHandshakeInterceptor implements HandshakeInterceptor {
    @Autowired
    private JwtServiceImpl jwtService;

    // 握手前检验token,认证成功后设置认证信息到attributes中
//    @Override
//    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
//        String authToken = request.getHeaders().getFirst("Authorization");
//
//        // 如果请求头中没有 Authorization，则尝试从 URL 查询参数中获取 token
//        if (authToken == null || authToken.isEmpty()) {
//            authToken = extractTokenFromQuery(request);
//        }
//
//        if (authToken == null || authToken.isEmpty()) {
//            response.setStatusCode(HttpStatus.UNAUTHORIZED);
//            return false; // 拒绝握手
//        }
//
//        try {
//            AuthDetails authDetails = jwtService.analyzeToken(authToken);
//            List<SimpleGrantedAuthority> authorities = authDetails.getAuthorities().stream()
//                    .map(auth -> new SimpleGrantedAuthority(auth.getAuthority()))
//                    .collect(Collectors.toList());
//
//            Authentication authentication = new UsernamePasswordAuthenticationToken(
//                    authDetails.getUsername(),
//                    null,
//                    authorities
//            );
//
//            attributes.put("user", new StompPrincipal(authDetails.getUsername(), authDetails.getUser().getId())); // 必须设置
//        } catch (Exception e) {
//            response.setStatusCode(HttpStatus.FORBIDDEN);
//            return false;
//        }
//
//        return true;
//    }


    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        String authToken = request.getHeaders().getFirst("Authorization");

        if (authToken == null || authToken.isEmpty()) {
            authToken = extractTokenFromQuery(request);
        }

        if (authToken == null || authToken.isEmpty()) {
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;
        }

        try {
            AuthDetails authDetails = jwtService.analyzeToken(authToken);
            List<SimpleGrantedAuthority> authorities = authDetails.getAuthorities().stream()
                    .map(auth -> new SimpleGrantedAuthority(auth.getAuthority()))
                    .collect(Collectors.toList());

            Authentication authentication = new UsernamePasswordAuthenticationToken(
                    authDetails.getUsername(),
                    null,
                    authorities
            );
            Principal principal = new StompPrincipal(authDetails.getUsername(), authDetails.getUser().getId());
            attributes.put("principal", principal );
            if (request instanceof ServletServerHttpRequest) {
                ((ServletServerHttpRequest) request).getServletRequest()
                        .setAttribute("principal", principal); // 供 STOMP 层使用
            }
        } catch (Exception e) {
            response.setStatusCode(HttpStatus.FORBIDDEN);
            return false;
        }

        return true;
    }
    /**
     * 从 URL 查询参数中提取 token
     */
    private String extractTokenFromQuery(ServerHttpRequest request) {
        String query = request.getURI().getQuery();
        if (query == null) {
            return null;
        }

        return Arrays.stream(query.split("&"))
                .filter(param -> param.startsWith("token="))
                .findFirst()
                .map(param -> param.replaceFirst("token=", ""))
                .orElse(null);
    }
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        Principal principal = request.getPrincipal();
        if (principal != null) {
            String sessionId = request.getHeaders().getFirst("Sec-WebSocket-Key"); // WebSocket session ID 近似值
            System.out.println("WebSocket 连接建立成功:");
            System.out.println("  用户名: " + principal.getName());
            System.out.println("  会话ID: " + (sessionId != null ? sessionId : "未知"));
            System.out.println("  Principal类型: " + principal.getClass().getSimpleName());
        } else {
            System.err.println("WebSocket 握手完成，但 Principal 为 null");
        }
    }
    // 自定义的 Principal 实现类
    public static class StompPrincipal implements Principal {
        private final String name;
        private final Integer userId;

        public StompPrincipal(String name, Integer userId) {
            this.name = name;
            this.userId = userId;
        }

        @Override
        public String getName() {
            return name;
        }

        public Integer getUserId() {
            return userId;
        }

        @Override
        public String toString() {
            return String.format("User: %s (ID: %d)", name, userId);
        }
    }
}
